from torch import nn
import torch
import torchvision
from torch.utils.data import DataLoader

dataset = torchvision.datasets.CIFAR10("./dataset_2",train=False,transform=torchvision.transforms.ToTensor(),download=True)
data_loader = DataLoader(dataset,1)

class MyNN(torch.nn.Module):
    def __init__(self):
        super(MyNN, self).__init__()
        self.model1 = nn.Sequential(nn.Conv2d(3,32,5,padding=2),
                                    nn.MaxPool2d(2),
                                    nn.Conv2d(32,32,5,padding=2),
                                    nn.MaxPool2d(2),
                                    nn.Conv2d(32,64,5,padding=2),
                                    nn.MaxPool2d(2),
                                    nn.Flatten(),
                                    nn.Linear(1024,64),
                                    nn.Linear(64,10))

    def forward(self,input):
        output = self.model1(input)
        return output

my_nn = MyNN()
loss = nn.CrossEntropyLoss()

for data in data_loader:
    imgs,target = data
    output = my_nn(imgs)
    # print(target)
    # print(output)
    res_loss = loss(output,target)
    # print(res_loss)
    # print(type(res_loss))
    res_loss.backward()
    print("ok")
